from mcp.server.fastmcp import FastMCP
import numpy as np

# 创建MCP服务器
mcp = FastMCP("Linear Algebra Matrix Operations Server")

@mcp.tool()
def matrix_addition(matrix_a: list, matrix_b: list) -> list:
    """Perform matrix addition"""
    return (np.array(matrix_a) + np.array(matrix_b)).tolist()

@mcp.tool()
def matrix_multiplication(matrix_a: list, matrix_b: list) -> list:
    """Perform matrix multiplication"""
    return np.matmul(matrix_a, matrix_b).tolist()

@mcp.tool()
def scalar_multiplication(matrix: list, scalar: float) -> list:
    """Multiply matrix by scalar"""
    return (np.array(matrix) * scalar).tolist()

@mcp.tool()
def matrix_transpose(matrix: list) -> list:
    """Calculate matrix transpose"""
    return np.array(matrix).T.tolist()

@mcp.tool()
def matrix_determinant(matrix: list) -> float:
    """Calculate matrix determinant"""
    return float(np.linalg.det(matrix))

@mcp.tool()
def matrix_inverse(matrix: list) -> list:
    """Calculate matrix inverse"""
    return np.linalg.inv(matrix).tolist()

@mcp.tool()
def matrix_rank(matrix: list) -> int:
    """Calculate matrix rank"""
    return int(np.linalg.matrix_rank(matrix))

@mcp.tool()
def eigenvalues_eigenvectors(matrix: list) -> dict:
    """Calculate eigenvalues and eigenvectors"""
    eigenvalues, eigenvectors = np.linalg.eig(matrix)
    return {
        'eigenvalues': eigenvalues.tolist(),
        'eigenvectors': eigenvectors.tolist()
    }

@mcp.tool()
def solve_linear_system(coefficients: list, constants: list) -> list:
    """Solve system of linear equations"""
    return np.linalg.solve(coefficients, constants).tolist()

if __name__ == "__main__":
    mcp.run(transport="stdio")
